import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm
import os
from PIL import Image
import platform
from utils import *
import sys
import io
import matplotlib


matplotlib.use("TkAgg")
# 检查设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")
# 应用字体设置
set_chinese_font()
# 查看训练集图片
visualize_samples()

data_path = r"D:\desktop\class\DataSet\cats_vs_dogs_labeled"

# 简化版数据变换（去掉数据增强，只保留基本变换）
simple_train_transform = transforms.Compose(
    [
        transforms.Resize((256, 256)),  # 调整大小
        transforms.ToTensor(),  # 转换为Tensor
        # 简化标准化（使用更简单的值）
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
    ]
)

simple_test_transform = transforms.Compose(
    [
        transforms.Resize((256, 256)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
    ]
)

# 使用简化变换重新创建数据集
train_dataset = ImageFolder(
    os.path.join(data_path, "train"), transform=simple_train_transform
)

test_dataset = ImageFolder(
    os.path.join(data_path, "test"), transform=simple_test_transform
)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)


class SimpleCatDogCNN(nn.Module):
    def __init__(self, num_classes=2):  # 输出2个节点
        super(SimpleCatDogCNN, self).__init__()
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(3 * 256 * 256, 128),
            nn.ReLU(inplace=True),
            nn.Linear(128, 2),  # 输出2个值
            # 没有激活函数，CrossEntropyLoss会处理
        )

    def forward(self, x):
        return self.classifier(x)


model = SimpleCatDogCNN(num_classes=2).to(device)
print("模型结构:")
print(model)


# 训练模型
def train_model(model, train_loader, test_loader, criterion, optimizer, num_epochs=10):
    train_losses = []
    test_losses = []
    train_accuracies = []
    test_accuracies = []

    best_test_acc = 0.0

    for epoch in range(num_epochs):
        print(f"\nEpoch {epoch+1}/{num_epochs}")
        print("-" * 50)

        # 训练阶段
        model.train()
        running_loss = 0.0
        running_corrects = 0

        # 使用进度条
        pbar = tqdm(train_loader, desc=f"训练")
        for inputs, labels in pbar:
            inputs = inputs.to(device)
            labels = labels.to(device)

            # 前向传播
            outputs = model(inputs)
            loss = criterion(outputs, labels)

            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # 统计
            running_loss += loss.item() * inputs.size(0)
            _, preds = torch.max(outputs, 1)
            running_corrects += torch.sum(preds == labels.data)

            # 更新进度条
            pbar.set_postfix(
                {
                    "Loss": f"{loss.item():.4f}",
                    "Acc": f"{torch.sum(preds == labels.data).item() / len(labels):.2f}",
                }
            )

        epoch_loss = running_loss / len(train_dataset)
        epoch_acc = running_corrects.double() / len(train_dataset)

        train_losses.append(epoch_loss)
        train_accuracies.append(epoch_acc.cpu().numpy())

        # 验证阶段
        model.eval()
        test_running_loss = 0.0
        test_running_corrects = 0

        with torch.no_grad():
            for inputs, labels in test_loader:
                inputs = inputs.to(device)
                labels = labels.to(device)

                outputs = model(inputs)
                loss = criterion(outputs, labels)

                test_running_loss += loss.item() * inputs.size(0)
                _, preds = torch.max(outputs, 1)
                test_running_corrects += torch.sum(preds == labels.data)

        test_epoch_loss = test_running_loss / len(test_dataset)
        test_epoch_acc = test_running_corrects.double() / len(test_dataset)

        test_losses.append(test_epoch_loss)
        test_accuracies.append(test_epoch_acc.cpu().numpy())

        print(f"训练损失: {epoch_loss:.4f}, 训练准确率: {epoch_acc:.4f}")
        print(f"验证损失: {test_epoch_loss:.4f}, 验证准确率: {test_epoch_acc:.4f}")

        # 保存最佳模型
        if test_epoch_acc > best_test_acc:
            best_test_acc = test_epoch_acc
            torch.save(model.state_dict(), "best_cat_dog_model.pth")
            print(f"新的最佳模型已保存! 验证准确率: {best_test_acc:.4f}")

    return train_losses, test_losses, train_accuracies, test_accuracies


# 可视化训练结果
def plot_training_results(train_losses, test_losses, train_accuracies, test_accuracies):
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))

    # 绘制损失曲线
    ax1.plot(train_losses, label="训练损失")
    ax1.plot(test_losses, label="验证损失")
    ax1.set_title("训练和验证损失")
    ax1.set_xlabel("Epoch")
    ax1.set_ylabel("Loss")
    ax1.legend()
    ax1.grid(True)

    # 绘制准确率曲线
    ax2.plot(train_accuracies, label="训练准确率")
    ax2.plot(test_accuracies, label="验证准确率")
    ax2.set_title("训练和验证准确率")
    ax2.set_xlabel("Epoch")
    ax2.set_ylabel("Accuracy")
    ax2.legend()
    ax2.grid(True)

    plt.tight_layout()
    plt.show()


criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

train_losses, test_losses, train_accuracies, test_accuracies = train_model(
    model, train_loader, test_loader, criterion, optimizer, num_epochs=10
)

# 绘制结果
plot_training_results(train_losses, test_losses, train_accuracies, test_accuracies)


# 测试模型预测
def predict_image(model, image_path, transform, class_names):
    """预测单张图片"""
    model.eval()

    # 加载图片
    image = Image.open(image_path).convert("RGB")
    image_tensor = transform(image).unsqueeze(0).to(device)

    with torch.no_grad():
        outputs = model(image_tensor)
        probabilities = torch.softmax(outputs, dim=1)
        confidence, predicted = torch.max(probabilities, 1)

        predicted_class = class_names[predicted.item()]
        confidence_testue = confidence.item()

    return predicted_class, confidence_testue


# 测试预测功能
def test_prediction():
    # 从验证集中随机选择一张图片测试
    import random

    random_index = random.randint(0, len(test_dataset) - 1)
    image_path, true_label = test_dataset.samples[random_index]
    true_class = test_dataset.classes[true_label]

    prediction, confidence = predict_image(
        model, image_path, simple_test_transform, test_dataset.classes
    )

    # 显示图片和预测结果
    image = Image.open(image_path).convert("RGB")
    plt.figure(figsize=(6, 6))
    plt.imshow(image)
    plt.title(f"真实: {true_class}\n预测: {prediction}\n置信度: {confidence:.2f}")
    plt.axis("off")
    plt.show()

    print(f"真实类别: {true_class}")
    print(f"预测类别: {prediction}")
    print(f"置信度: {confidence:.2f}")


# 测试预测
test_prediction()
